import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# 结构相似性损失（SSIM），用于衡量预测与标签的结构相似度，提升结构还原能力
class SSIM(nn.Module):
    def __init__(self, window_size=11, size_average=True):
        super(SSIM, self).__init__()
        self.window_size = window_size
        self.size_average = size_average
        self.channel = 1
        self.window = self.create_window(window_size)

    # 生成高斯核
    def gaussian(self, window_size, sigma):
        gauss = torch.Tensor([math.exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
        return gauss/gauss.sum()

    # 创建二维高斯窗口
    def create_window(self, window_size):
        _1D_window = self.gaussian(window_size, 1.5).unsqueeze(1)
        _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
        window = _2D_window.expand(1, 1, window_size, window_size)
        return window

    def forward(self, img1, img2):
        try:
            (_, channel, _, _) = img1.size()
            # 若通道数或类型变化则重新生成window
            if channel == self.channel and self.window.data.type() == img1.data.type():
                window = self.window
            else:
                window = self.create_window(self.window_size)
                window = window.to(img1.device)
                self.window = window
                self.channel = channel

            # 计算均值
            mu1 = F.conv2d(img1, window, padding=self.window_size//2, groups=channel)
            mu2 = F.conv2d(img2, window, padding=self.window_size//2, groups=channel)
            mu1_sq = mu1.pow(2)
            mu2_sq = mu2.pow(2)
            mu1_mu2 = mu1 * mu2

            # 计算方差和协方差
            sigma1_sq = F.conv2d(img1 * img1, window, padding=self.window_size//2, groups=channel) - mu1_sq
            sigma2_sq = F.conv2d(img2 * img2, window, padding=self.window_size//2, groups=channel) - mu2_sq
            sigma12 = F.conv2d(img1 * img2, window, padding=self.window_size//2, groups=channel) - mu1_mu2

            # SSIM公式中的常数
            C1 = (0.01) ** 2
            C2 = (0.03) ** 2

            numerator = (2 * mu1_mu2 + C1) * (2 * sigma12 + C2)
            denominator = (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
            denominator = torch.clamp(denominator, min=1e-8)  # 防止除零
            ssim_map = numerator / denominator

            # 返回1-ssim，作为损失
            if self.size_average:
                return torch.clamp(1 - ssim_map.mean(), min=0.0, max=1.0)
            else:
                return torch.clamp(1 - ssim_map.mean(1).mean(1).mean(1), min=0.0, max=1.0)
        except Exception as e:
            print(f"SSIM Error: {str(e)}")
            return torch.tensor(0.0, device=img1.device)

# 边缘检测损失，利用Sobel算子提取边缘，计算预测与标签边缘的L1损失
class EdgeDetectionLoss(nn.Module):
    def __init__(self):
        super(EdgeDetectionLoss, self).__init__()
        self.sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).reshape(1, 1, 3, 3)
        self.sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).reshape(1, 1, 3, 3)
        
    def forward(self, pred, target):
        try:
            # 将Sobel算子放到正确的设备
            if not hasattr(self, 'sobel_x_cuda') or self.sobel_x_cuda.device != pred.device:
                self.sobel_x_cuda = self.sobel_x.to(pred.device)
                self.sobel_y_cuda = self.sobel_y.to(pred.device)

            pred = torch.sigmoid(torch.clamp(pred, min=-10, max=10))
            
            # 计算预测和标签的边缘图
            pred_x = F.conv2d(pred, self.sobel_x_cuda, padding=1)
            pred_y = F.conv2d(pred, self.sobel_y_cuda, padding=1)
            pred_edge = torch.sqrt(pred_x.pow(2) + pred_y.pow(2) + 1e-8)

            target_x = F.conv2d(target, self.sobel_x_cuda, padding=1)
            target_y = F.conv2d(target, self.sobel_y_cuda, padding=1)
            target_edge = torch.sqrt(target_x.pow(2) + target_y.pow(2) + 1e-8)

            edge_loss = F.l1_loss(pred_edge, target_edge)
            return torch.clamp(edge_loss, min=0.0, max=1.0)
        except Exception as e:
            print(f"Edge Detection Error: {str(e)}")
            return torch.tensor(0.0, device=pred.device)

# 加权BCE和IoU损失，兼顾像素级和区域级的准确性
class WeightedBCEAndIoULoss(nn.Module):
    def __init__(self, smooth=1e-6):
        super(WeightedBCEAndIoULoss, self).__init__()
        self.smooth = smooth
        
    def forward(self, pred, target):
        try:
            pred = torch.sigmoid(torch.clamp(pred, min=-10, max=10))
            
            # 计算正负样本权重，缓解类别不平衡
            total_pixels = float(target.size(0) * target.size(2) * target.size(3))
            pos_weight = torch.clamp((target.sum() + self.smooth) / total_pixels, min=0.1, max=0.9)
            neg_weight = 1 - pos_weight
            
            # 加权BCE损失
            pred_clipped = torch.clamp(pred, min=self.smooth, max=1-self.smooth)
            bce_pos = -pos_weight * target * torch.log(pred_clipped)
            bce_neg = -neg_weight * (1 - target) * torch.log(1 - pred_clipped)
            bce_loss = torch.clamp((bce_pos + bce_neg).mean(), min=0.0, max=1.0)
            
            # IoU损失
            intersection = (pred * target).sum()
            union = pred.sum() + target.sum() - intersection + self.smooth
            iou = torch.clamp(intersection / union, min=self.smooth)
            iou_loss = 1 - iou
            
            # 动态调整BCE和IoU的权重
            if iou.item() < 0.5:
                total_loss = 0.4 * bce_loss + 0.6 * iou_loss
            else:
                total_loss = 0.6 * bce_loss + 0.4 * iou_loss
                
            return torch.clamp(total_loss, min=0.0, max=1.0)
        except Exception as e:
            print(f"BCE-IoU Error: {str(e)}")
            return torch.tensor(0.0, device=pred.device)

# 自适应Focal Loss，提升对难分样本的关注
class AdaptiveFocalLoss(nn.Module):
    def __init__(self, gamma=1.0, alpha=0.25, smooth=1e-6):
        super(AdaptiveFocalLoss, self).__init__()
        self.gamma = gamma
        self.alpha = alpha
        self.smooth = smooth
        
    def forward(self, pred, target):
        try:
            pred = torch.sigmoid(torch.clamp(pred, min=-10, max=10))
            pred = torch.clamp(pred, min=self.smooth, max=1-self.smooth)
            
            pt = target * pred + (1 - target) * (1 - pred)
            adaptive_gamma = torch.clamp(self.gamma * (1 - pt.mean()), min=0.0, max=2.0)
            
            focal_weight = (1 - pt) ** adaptive_gamma
            bce_loss = -target * torch.log(pred) - (1 - target) * torch.log(1 - pred)
            focal_loss = self.alpha * focal_weight * bce_loss
            
            return torch.clamp(focal_loss.mean(), min=0.0, max=1.0)
        except Exception as e:
            print(f"Focal Loss Error: {str(e)}")
            return torch.tensor(0.0, device=pred.device)

# 综合显著性损失，将多种损失加权融合
class ImprovedSaliencyLoss(nn.Module):
    def __init__(self):
        super(ImprovedSaliencyLoss, self).__init__()
        self.bce_iou_loss = WeightedBCEAndIoULoss()
        self.focal_loss = AdaptiveFocalLoss()
        self.ssim_loss = SSIM()
        self.edge_loss = EdgeDetectionLoss()
        
    def forward(self, pred, target):
        try:
            # 限制输入范围
            pred = torch.clamp(pred, min=-10, max=10)
            target = torch.clamp(target, min=0, max=1)
            
            # 计算各项损失
            bce_iou = self.bce_iou_loss(pred, target)
            focal = self.focal_loss(pred, target)
            ssim = self.ssim_loss(torch.sigmoid(pred), target)
            edge = self.edge_loss(pred, target)
            
            # 检查NaN，若有则用0.1替代
            if torch.isnan(bce_iou):
                print("Warning: BCE-IoU loss is NaN")
                bce_iou = torch.tensor(0.1, device=pred.device)
            if torch.isnan(focal):
                print("Warning: Focal loss is NaN")
                focal = torch.tensor(0.1, device=pred.device)
            if torch.isnan(ssim):
                print("Warning: SSIM loss is NaN")
                ssim = torch.tensor(0.1, device=pred.device)
            if torch.isnan(edge):
                print("Warning: Edge loss is NaN")
                edge = torch.tensor(0.1, device=pred.device)
            
            # 多损失加权融合
            total_loss = (0.4 * bce_iou + 
                         0.2 * focal + 
                         0.25 * ssim + 
                         0.15 * edge)
            
            return torch.clamp(total_loss, min=0.0, max=1.0)
        except Exception as e:
            print(f"Total Loss Error: {str(e)}")
            return torch.tensor(0.1, device=pred.device)  # 返回一个小的非零值而不是0